## special lattice support version
## method有'full'默认,以及'half_hermitian','half_symmetric',对付不同的关联矩阵
## 添加自定义pairs的关联测量数据处理，method用'custom'，返回矩阵中没有测量的全部赋NaN
def get_correlation_matrices_v4(propstring, correlator_string, method='full',special_lattice=0):
    import numpy as np
    def get_correlation_matrix_of_dataset(dataset, method):
        L = int(dataset.props['L'])
        correlation_order = dataset.x
        correlation_data = dataset.y
        ## 不同于dmrg，mps_optim算出的correlator是乱序的，还缺少主对角线上的元素
        correlation_matrix = np.zeros((L, L))
        correlation_matrix = correlation_matrix.astype(complex)  ## to support complex
        correlation_matrix[:] = complex('NaN')
        if method == 'custom':
            if len(np.shape(correlation_order))==1:
                measure_input = dataset.props['MEASURE_LOCAL_AT['+dataset.props['observable']+']']
                measure_pairs = measure_input.split('|')[1]
                measure_pairs_split = np.array([i for i in measure_pairs])
                b0 = np.where(measure_pairs_split=='(')
                b1 = np.where(measure_pairs_split==',')
                idx1 = b0[0][0]
                idx2 = b1[0][0]
                rowstring = measure_pairs[idx1+1:idx2]
                row = int(rowstring)
                correlation_order = np.array([[row,i] for i in correlation_order])
        n = 0
        if special_lattice==0:
            for order in correlation_order:
                correlation_matrix[order[0], order[1]] = correlation_data[0][n]
                n += 1
        elif special_lattice==1:
            for order in correlation_order:
                correlation_matrix[order[0]-1, order[1]-1] = correlation_data[0][n]
                n += 1
        for i in range(L):
            correlation_matrix[i, i] = complex('NaN')
        calcued_elements = np.shape(correlation_order)[0]
        if method == 'full':
            assert calcued_elements == L**2 - L, 'Report: Error, you put half measured correlation matrix into full one program'
        elif method == 'half_hermitian':
            assert calcued_elements != L**2 - L, 'Report: Error, you put full measured correlation matrix into half one program'
            correlation_matrix = correlation_matrix + correlation_matrix.conj().T
        elif method == 'half_symmetric':
            assert calcued_elements != L**2 - L, 'Report: Error, you put full measured correlation matrix into half one program'
            correlation_matrix = correlation_matrix + correlation_matrix.T
        elif method =='custom':
            return correlation_matrix
        else:
            raise Exception, "Invalid method for my correlation matrix !"
        return correlation_matrix
    eigen_measure_obs = pyalps.loadEigenstateMeasurements(
        result_files, what=[correlator_string])
    ku = [i[0] for i in eigen_measure_obs]
    def by_props(dataset):  ##把DataSet按照其props中的mu值排序
        return dataset.props[propstring]
    ## 自己实现collectXY就比较灵活
    obs_sorted = sorted(ku, key=by_props)
    paras = [i.props[propstring] for i in obs_sorted]
    correlation_matrices = [get_correlation_matrix_of_dataset(i, method) for i in obs_sorted]
    paras = np.array(paras)
    correlation_matrices = np.array(correlation_matrices)
    return (paras, correlation_matrices)


(Ds, SzSz_correlation_matrices) = get_correlation_matrices_v4('MAXSTATES',
 'custom_SzSz', 'custom',special_lattice=1)
